import torch
import torch.nn as nn
import torch.nn.functional as F
import math

"""
由于原始代码依赖于华为 NPU 的特定算子，我将在标准 PyTorch 环境中创建一个简化但功能完整的演示，展示 Prefill 阶段的核心注意力计算逻辑。
1. RotaryPositionEmbedding 类
实现了旋转位置编码 (RoPE) 的核心算法

通过正弦和余弦函数为不同位置的 token 生成不同的编码

2. mock_npu_fused_infer_attention_score 函数
模拟了华为 NPU 的 npu_fused_infer_attention_score 算子的功能

实现了完整的注意力计算流程：

应用旋转位置编码（如果提供）

计算 QK^T 注意力分数

应用注意力掩码（包括因果掩码）

计算 Softmax 注意力权重

应用注意力权重到值向量

3. SimplifiedDeepseekMLA 类
简化的 DeepseekMLA 模型实现

包含所有必要的投影层和归一化层

forward_prefill 方法实现了完整的 Prefill 阶段前向传播

4. 演示功能
创建模拟输入数据

运行 Prefill 阶段计算

可视化注意力权重，展示模型如何关注输入序列中的不同部分

这个演示代码可以在标准 PyTorch 环境中运行，不需要华为 NPU 硬件或软件栈。它展示了 Prefill 阶段的核心计算逻辑，包括旋转位置编码的应用和注意力机制的计算。
"""


class RotaryPositionEmbedding(nn.Module):
    """旋转位置编码 (RoPE) 的简化实现"""

    def __init__(self, dim, max_seq_len=2048, base=10000):
        super().__init__()
        self.dim = dim
        self.max_seq_len = max_seq_len
        self.base = base

        # 计算频率
        inv_freq = 1.0 / (base ** (torch.arange(0, dim, 2).float() / dim))
        self.register_buffer('inv_freq', inv_freq)

    def forward(self, x, positions=None):
        """
        x: [batch_size, seq_len, num_heads, head_dim]
        positions: [batch_size, seq_len]
        """
        batch_size, seq_len, num_heads, head_dim = x.shape

        if positions is None:
            positions = torch.arange(seq_len, device=x.device).unsqueeze(0).expand(batch_size, seq_len)

        # 生成正弦和余弦
        freqs = torch.einsum('bi,j->bij', positions.float(), self.inv_freq)  # [batch_size, seq_len, dim/2]
        freqs = torch.cat([freqs, freqs], dim=-1)  # [batch_size, seq_len, dim]

        # 重塑为与x相同的形状
        cos = torch.cos(freqs).unsqueeze(2).expand_as(x)  # [batch_size, seq_len, num_heads, head_dim]
        sin = torch.sin(freqs).unsqueeze(2).expand_as(x)  # [batch_size, seq_len, num_heads, head_dim]

        # 应用旋转位置编码
        x1, x2 = x[..., :self.dim // 2], x[..., self.dim // 2:]
        rx = torch.cat([-x2, x1], dim=-1)

        return x * cos + rx * sin


def mock_npu_fused_infer_attention_score(
        q, k, v, query_rope=None, key_rope=None, num_heads=1,
        num_key_value_heads=1, input_layout="TND", atten_mask=None,
        sparse_mode=0, actual_seq_lengths=None, actual_seq_lengths_kv=None, attn_mask=None,
        scale=1.0, next_tokens=0
):
    """
    模拟华为 NPU 的 npu_fused_infer_attention_score 算子
    这个函数在标准 PyTorch 环境中实现相同的功能
    """
    # 确保输入形状正确
    assert input_layout == "TND", "只支持 TND 布局"

    # 应用旋转位置编码 (如果提供了)
    if query_rope is not None:
        q = q + query_rope
    if key_rope is not None:
        k = k + key_rope

    # 计算注意力分数: Q * K^T
    attn_scores = torch.matmul(q, k.transpose(-2, -1)) * scale

    # 应用注意力掩码
    if atten_mask is not None:
        if sparse_mode == 3:  # 因果掩码
            # 创建因果掩码
            seq_len = q.size(-2)
            causal_mask = torch.tril(torch.ones(seq_len, seq_len, device=q.device)).bool()
            attn_scores = attn_scores.masked_fill(~causal_mask, float('-inf'))
        else:
            # 应用提供的掩码
            attn_scores = attn_scores + attn_mask

    # 计算注意力权重
    attn_weights = F.softmax(attn_scores, dim=-1)

    # 应用注意力权重到值
    attn_output = torch.matmul(attn_weights, v)

    return attn_output, attn_weights


class SimplifiedDeepseekMLA(nn.Module):
    """简化的 DeepseekMLA Prefill 阶段实现"""

    def __init__(self, hidden_size=4096, num_heads=32, qk_nope_head_dim=48,
                 qk_rope_head_dim=16, v_head_dim=64, q_lora_rank=64, kv_lora_rank=64):
        super().__init__()
        self.hidden_size = hidden_size
        self.num_heads = num_heads
        self.qk_nope_head_dim = qk_nope_head_dim
        self.qk_rope_head_dim = qk_rope_head_dim
        self.qk_head_dim = qk_nope_head_dim + qk_rope_head_dim
        self.v_head_dim = v_head_dim
        self.q_lora_rank = q_lora_rank
        self.kv_lora_rank = kv_lora_rank

        # 模拟的投影层
        self.q_proj = nn.Linear(hidden_size, num_heads * self.qk_head_dim)
        self.kv_proj = nn.Linear(hidden_size, kv_lora_rank + qk_rope_head_dim)
        self.kv_b_proj = nn.Linear(kv_lora_rank, num_heads * (qk_nope_head_dim + v_head_dim))
        self.o_proj = nn.Linear(num_heads * v_head_dim, hidden_size)

        # 模拟的层归一化
        self.kv_layernorm = nn.LayerNorm(kv_lora_rank)

        # 旋转位置编码
        self.rotary_emb = RotaryPositionEmbedding(qk_rope_head_dim)

        # 缩放因子
        self.scale = self.qk_head_dim ** -0.5

        # 注意力掩码
        self.attn_mask = ~torch.tril(torch.ones((2048, 2048), dtype=torch.bool))

    def forward_prefill(self, hidden_states, positions):
        """简化的 Prefill 前向传播"""
        batch_size, seq_len, _ = hidden_states.shape

        # 1. 查询投影
        q = self.q_proj(hidden_states)
        q = q.view(batch_size, seq_len, self.num_heads, self.qk_head_dim)

        # 2. 键值投影
        kv_latent = self.kv_proj(hidden_states)
        kv_a, k_pe_raw = torch.split(kv_latent, [self.kv_lora_rank, self.qk_rope_head_dim], dim=-1)
        kv_a = self.kv_layernorm(kv_a)

        # 3. 应用旋转位置编码
        # cos, sin = self.rotary_emb.get_cos_sin(positions)

        # 分割查询
        q_nope, q_pe = torch.split(q, [self.qk_nope_head_dim, self.qk_rope_head_dim], dim=-1)

        # 应用 RoPE 到查询
        q_pe_rotated = self.rotary_emb(q_pe, positions)
        q = torch.cat([q_nope, q_pe_rotated], dim=-1)

        # 应用 RoPE 到键
        k_pe_rotated = self.rotary_emb(
            k_pe_raw.view(batch_size, seq_len, 1, self.qk_rope_head_dim),
            positions
        ).squeeze(2)

        # 4. 键值最终投影
        kv = self.kv_b_proj(kv_a)
        kv = kv.view(batch_size, seq_len, self.num_heads, self.qk_nope_head_dim + self.v_head_dim)
        k_nope, v = torch.split(kv, [self.qk_nope_head_dim, self.v_head_dim], dim=-1)

        # 合并键
        k = torch.cat([k_nope, k_pe_rotated.unsqueeze(2).repeat(1, 1, self.num_heads, 1)], dim=-1)

        # 5. 调整形状以适应注意力计算
        q = q.transpose(1, 2)  # [batch_size, num_heads, seq_len, head_dim]
        k = k.transpose(1, 2)
        v = v.transpose(1, 2)

        # 6. 使用模拟的注意力算子
        attn_output, attn_weights = mock_npu_fused_infer_attention_score(
            q, k, v,
            query_rope=None,  # 已经应用了 RoPE
            key_rope=None,  # 已经应用了 RoPE
            num_heads=self.num_heads,
            num_key_value_heads=self.num_heads,
            input_layout="TND",
            atten_mask=self.attn_mask[:seq_len, :seq_len],
            sparse_mode=3,  # 因果掩码
            scale=self.scale,
            next_tokens=0
        )

        # 7. 调整形状并应用输出投影
        attn_output = attn_output.transpose(1, 2).contiguous()
        attn_output = attn_output.view(batch_size, seq_len, -1)
        output = self.o_proj(attn_output)

        return output, attn_weights


# 演示代码
def demo_prefill_attention():
    """演示 Prefill 阶段的注意力计算"""
    # 设置参数
    batch_size = 2
    seq_len = 10
    hidden_size = 4096

    # 创建模型
    model = SimplifiedDeepseekMLA()

    # 创建模拟输入
    hidden_states = torch.randn(batch_size, seq_len, hidden_size)
    positions = torch.arange(seq_len).unsqueeze(0).expand(batch_size, seq_len)

    # 运行 Prefill
    output, attn_weights = model.forward_prefill(hidden_states, positions)

    print("输入形状:", hidden_states.shape)
    print("输出形状:", output.shape)
    print("注意力权重形状:", attn_weights.shape)

    # 可视化注意力权重（第一个样本，第一个头）
    import matplotlib.pyplot as plt

    plt.figure(figsize=(10, 8))
    plt.imshow(attn_weights[0, 0].detach().numpy(), cmap='viridis')
    plt.colorbar()
    plt.title("注意力权重 (第一个样本，第一个头)")
    plt.xlabel("Key Position")
    plt.ylabel("Query Position")
    plt.show()

    return output, attn_weights


if __name__ == "__main__":
    output, attn_weights = demo_prefill_attention()
